/*
 * @Description  : 利用MINST训练KNN模型，并识别图片中的数字
 * @Author       : hx
 * @Date         : 2021-04-11 12:00:01
 * @LastEditors: Gtylcara.
 * @LastEditTime: 2021-04-24 22:47:50
 * @FilePath     : /minst/src/main.cpp
 */
#include <fstream>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <string>
using namespace std;
using namespace cv;

// 小端存储转换
int ReverseInt(int data);
// 读取image数据集信息
Mat ReadMINSTImgs(const String PATH);
// 读取label数据集信息
Mat ReadMINSTLabs(const String PATH);
// 生成KNN模型并利用MINST数据集测试模型准确率
void TrainKNNModule(Mat trainImgs, Mat trainLabs, Mat testImgs, Mat testLabs);
// 利用KNN模型识别图片中的数字
void ImagePredict(void);

const String TRAIN_IMGS = "../img/train-images.idx3-ubyte";
const String TRAIN_LABS = "../img/train-labels.idx1-ubyte";
const String TEST_IMGS = "../img/t10k-images.idx3-ubyte";
const String TEST_LABS = "../img/t10k-labels.idx1-ubyte";
const String MODEL_PATH = "../model/mnist_knn.xml";
const String PIC_PATH = "../img/pic.jpg";

int main() {
  /*
  ---------读取训练数据-----------
  */
  Mat trainLabs = ReadMINSTLabs(TRAIN_LABS);
  // 读取训练标签数据 (60000,1) 类型为int32
  Mat trainImgs = ReadMINSTImgs(TRAIN_IMGS);
  // 读取训练图像数据 (60000,784) 类型为float32 数据未归一化
  trainImgs = trainImgs / 255.0;
  // 将图像数据归一化

  /*
  ---------读取测试数据-----------
  */
  Mat testLabs = ReadMINSTLabs(TEST_LABS);
  // 读取测试数据标签(10000,1) 类型为int32
  Mat testImgs = ReadMINSTImgs(TEST_IMGS);
  // 读取测试数据图像 (10000,784) 类型为float32 数据未归一化
  testImgs = testImgs / 255.0;
  // 归一化

  /*
  ---------KNN模型构建与识别-----------
  */
  TrainKNNModule(trainImgs, trainLabs, testImgs, testLabs);
  // 构建并测试KNN模型
  ImagePredict();
  // 预测图片中的数字

  return 0;
}

/**
 * @description: 大端存储转小端存储
 * @param {int} data
 * @return {int}
 */
int ReverseInt(int data) {
  unsigned char high16, high8, low16, low8;

  high16 = data & 255;
  // 获取数据的二进制的1～8位
  high8 = (data >> 8) & 255;
  // 获取数据的二进制的9～16位
  low16 = (data >> 16) & 255;
  // 获取数据的二进制的17～24位
  low8 = (data >> 24) & 255;
  // 获取数据的二进制的25～32位

  return ((int)high16 << 24) + ((int)high8 << 16) + ((int)low16 << 8) + low8;
  // 将high16移到25~32位，将high8移到17~24位，将low16移到9~16位，将low8移到1~8位
}

/**
 * @description: 读取MINST样本的图像
 * @param {const String} PATH
 * @return {Mat} imageMat
 */
Mat ReadMINSTImgs(const String PATH) {
  Mat imageMat;
  int magicNum, imgNum, imgRows, imgCols;

  ifstream file(PATH, ios::binary);
  // 以二进制的形式读取MINST数据集
  if (true == file.fail()) {
    // 判断是否读取失败
    cout << "MINST训练图像读取失败" << endl;
    exit(-1);
  }

  file.read((char *)&magicNum, sizeof(magicNum));
  // 读取数据集中图像的序号
  magicNum = ReverseInt(magicNum);
  // 转换序号的存储方式（大端存储->小端存储）
  file.read((char *)&imgNum, sizeof(imgNum));
  // 读取数据集中图像的数量
  imgNum = ReverseInt(imgNum);
  file.read((char *)&imgRows, sizeof(imgRows));
  // 读取数据集中图像的高度
  imgRows = ReverseInt(imgRows);
  file.read((char *)&imgCols, sizeof(imgCols));
  // 读取数据集中图像的宽度
  imgCols = ReverseInt(imgCols);
  cout << "magicNum:" << magicNum << ' ' << "imgNum:" << imgNum << " "
       << "imgRows:" << imgRows << " "
       << "imgCols:" << imgCols << endl;

  cout << "开始读取Image数据......" << endl;
  imageMat = Mat::zeros(imgNum, imgRows * imgCols, CV_32FC1);
  // 将imageMat初始化为imgNum个图像大小（像素点数量）为imgRows * imgCols的集合
  // 而且集合中每个图片的类型为float32
  for (int i = 0; i < imgNum; i++) {
    // 遍历数据集中的每一个图像
    for (int j = 0; j < imgRows * imgCols; j++) {
      // 遍历图像中的每一个像素点
      unsigned char temp = 0;
      file.read((char *)&temp, sizeof(temp));
      // 读取数据集的图像中的每一个像素点
      float pixel_value = float(temp);
      // 将每个像素值的类型转换成float型
      imageMat.at<float>(i, j) = pixel_value;
      // 按照行将像素值一个个写入Mat中
    }
  }
  cout << "读取Image数据完毕......" << endl;
  file0.close();
  file.close();
  // 关闭文件
  return imageMat;
}

/**
 * @description: 读取MINST样本的标签
 * @param {const String} PATH
 * @return {Mat} lableMat
 */
Mat ReadMINSTLabs(const String PATH) {
  Mat lableMat;
  int magicNum, labNum;
  ifstream file(PATH, ios::binary);
  if (true == file.fail()) {
    cout << "MINST训练标签读取失败" << endl;
    exit(-1);
  }

  file.read((char *)&magicNum, sizeof(magicNum));
  magicNum = ReverseInt(magicNum);
  file.read((char *)&labNum, sizeof(labNum));
  labNum = ReverseInt(labNum);
  cout << "magicNum:" << magicNum << ' ' << "labNum:" << labNum << endl;

  cout << "开始读取Label数据......" << endl;
  lableMat = Mat::zeros(labNum, 1, CV_32SC1);
  // CV_32SC1代表32位有符号整型 通道数为1
  for (int i = 0; i < labNum; i++) {
    unsigned char temp = 0;
    file.read((char *)&temp, sizeof(temp));
    lableMat.at<unsigned int>(i, 0) = (unsigned int)temp;
  }
  cout << "读取Label数据完毕......" << endl;

  file.close();
  return lableMat;
}

/**
 * @description: 构建KNN模型并测试模型
 * @param {Mat} trainImgs
 * @param {Mat} trainLabs
 * @param {Mat} testImgs
 * @param {Mat} testLabs
 * @return {*}
 */
void TrainKNNModule(Mat trainImgs, Mat trainLabs, Mat testImgs, Mat testLabs) {
  /*
  ---------通过MINST训练集构建KNN模型-----------
  */
  cv::Ptr<cv::ml::KNearest> knn = cv::ml::KNearest::create();
  // 声明KNN
  knn->setDefaultK(5);
  // 设置K值 默认的k为10
  knn->setIsClassifier(true);
  // 设置分类还是回归 这里是分类

  cout << "开始进行训练..." << endl;
  cv::Ptr<cv::ml::TrainData> train_data =
      cv::ml::TrainData::create(trainImgs, cv::ml::ROW_SAMPLE, trainLabs);
  // 装载训练数据和配置训练方法
  knn->train(train_data);
  // 开始训练
  cout << "训练完成" << endl;

  /*
  ---------利用MINST测试集验证KNN模型-----------
  */
  Mat resultMat;
  cout << "开始进行预测..." << endl;
  knn->predict(testImgs, resultMat);
  // testImgs 是测试集的图像，resultMat 是预测结果的集合
  cout << "预测完成" << endl;

  resultMat.convertTo(resultMat, CV_8UC1);
  testLabs.convertTo(testLabs, CV_8UC1);
  // resultMat需要与testLabs比较，所以resultMat需要转换成与testLabs一样的类型

  int passCnt = 0;
  // 识别正确的数量
  for (int i = 0; i < resultMat.rows; i++) {
    if (resultMat.at<uchar>(i, 0) == testLabs.at<uchar>(i, 0)) {
      // 将KNN识别出的数字与MINST标记的数据进行对比
      passCnt++;
    }
  }
  float res = float(passCnt) / float(resultMat.rows);
  // 正确识别的数量 / 识别的总数
  cout << "测试数据集上的准确率为：" << res * 100 << "%" << endl;

  knn->save(MODEL_PATH);
  // 保存模型
}

void ImagePredict(void) {
  /*
  ---------读取图片，并对其进行预处理(KNN要求的)-----------
  */
  cout << "开始读取图片" << endl;
  Mat image = cv::imread(PIC_PATH, 0);
  // 读取图片
  Mat img_show = image.clone();
  // 备份原图
  resize(image, image, Size(28, 28));
  // 将图片裁剪成训练集中图片的大小
  image.convertTo(image, CV_32F);
  // 将image的类型转换为float32
  image = image / 255.0;
  // 将图片归一化
  image = image.reshape(1, 1);
  // 将图片序列化
  cout << "图片处理完成" << endl;

  /*
  ---------利用KNN模型进行识别-----------
  */
  cout << "开始载入KNN模型" << endl;
  cv::Ptr<cv::ml::KNearest> knn =
      cv::ml::StatModel::load<cv::ml::KNearest>(MODEL_PATH);
  // 加载knn模型
  cout << "开始识别图片" << endl;
  float res = knn->predict(image);
  // 识别图片
  cout << "图片中的数字是" << res << endl;
  cv::imshow("img", img_show);
  // 显示图片
  cv::waitKey(0);
  // 等待按下任意键
}
